/*
 * Copyright 2008 University Corporation for Advanced Internet Development, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
/*
 * Additional modifications by Clareity Security, LLC for use as a DB backed storage service.
 * 
 */
package net.clareitysecurity.shibboleth.storage;

import java.util.Iterator;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;

import net.clareitysecurity.shibboleth.db.SessionObjectDAOImpl;

import org.joda.time.DateTime;
import org.joda.time.Interval;
import org.joda.time.chrono.ISOChronology;
import org.opensaml.util.storage.StorageService;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.context.ApplicationContext;
import org.springframework.context.ApplicationContextAware;
import org.springframework.context.ApplicationEvent;


import edu.internet2.middleware.shibboleth.idp.session.impl.SessionManagerEntry;

/**
 * A thread-safe implementation of {@link StorageService} that publishes event when items are added or removed from the
 * service.
 * 
 * An {@link AddEntryEvent} is published after an item has been added to the storage service. A {@link RemoveEntryEvent}
 * is published after an item has been removed from the storage service. These events are published in the root
 * application context, that is the highest ancestor, of the application context presented to this class.
 * 
 * @param <KeyType> object type of the keys
 * @param <ValueType> object type of the values
 */
public class DbStorageService<KeyType, ValueType> implements StorageService<KeyType, ValueType>,
        ApplicationContextAware {

	private final Logger log = LoggerFactory.getLogger(DbStorageService.class);
	public final static String version = "Version 1.0.3 (2013-08-27)";

    /** Spring application context. */
    private ApplicationContext appCtx;

    /** Backing map. */
    private Map<String, Map<KeyType, ValueType>> store;
    // the map for the sessions
    private Map<KeyType, ValueType> sessionCache;
    
    // The special name of the session partition
    public final static String sessionPartition  = "session";

    private String systemName = null;
    
    private StorageCleaner<KeyType> storageCleaner = null;
    private Thread storageThread = null;
    
    /**
     * Assumptions being made about the Shibboleth storage service:
     * 
     * The partition containing session information is called "session". That is the only
     * partition that we will store in the Infinispan backed cache. Other partitions
     * in use include:
     * 
     *   loginContexts -- contains information during the authentication cycle
     *   replay        -- contains SAML ID values used to detect replay attacks
     *   
     * These and any other partitions will be stored locally only.
     */
    
    /** Constructor. */
    public DbStorageService(String systemName) throws Exception {
    	
    	try {
    		
            String endofline = System.getProperty("line.separator");
            log.info(endofline + "=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=" + endofline + endofline);

	    	log.info(version);
	        store = new ConcurrentHashMap<String, Map<KeyType, ValueType>>();
	        // explicitly add the session partition
	        sessionCache = new ConcurrentHashMap<KeyType, ValueType>();
	        store.put(sessionPartition, sessionCache);

	        this.systemName = systemName;
	        log.info("Using systemName of [{}]", systemName);
	        
	        log.info(endofline);
	        
	        // Create the Storage Cleaner object
	        storageCleaner = new StorageCleaner<KeyType>();
	        storageCleaner.setSystemName(systemName);
	        storageCleaner.setWork(true);
	        storageCleaner.setStorageService(this);
		    // Create the thread supplying it with the runnable object
		    storageThread = new Thread(storageCleaner);
		    storageThread.start();

    	} catch (Exception e) {
    		log.error("Unable to instantiate ShibDB cache. Exception is [{}]", e.getMessage());
    		log.error("{}", e);
    		
    		throw new Exception("Unable to instantiate ShibDB cache.", e);
    	}
    }

    /**
     * The ClusterFilter class is coded to call this when it is unloaded by the application engine.
     */
    public void cleanup() {
    	log.info("Cleaning up cache resources.");
    	if (storageCleaner != null) {
    		storageCleaner.setWork(false);
        	try {
        		storageThread.interrupt(); // it's probably asleep, wake it up
        		Thread.sleep(250);         // yield execution
        	} catch (Exception e) {}
    	}
    }
    
    public String checkStatus(boolean fullStorageCheck) {
    	try {
	    	DateTime now = new DateTime(ISOChronology.getInstanceUTC());
	    	StringBuilder sb = new StringBuilder();
	    	sb.append("Checking session cache information ...\r\n");
	    	
	    	sb.append("Cache size = [" + sessionCache.size() + "]\r\n");
	    	
	    	if (fullStorageCheck == true) {
	    		sb.append("Items in cache:\r\n");
	    		Set<KeyType> keys = sessionCache.keySet();
	    		for (Object o: keys) {
	    			sb.append("    [" + o.toString() + "] = " + sessionCache.get(o).getClass().toString() + "\r\n");
	    			if (sessionCache.get(o) instanceof SessionManagerEntry) {
	    				SessionManagerEntry sme = (SessionManagerEntry) sessionCache.get(o);
	    				long minutes = 0;
	    				try {
		    				if (now.isBefore(sme.getExpirationTime().getMillis()) == true) {
		    					minutes = new Interval(now, sme.getExpirationTime()).toDurationMillis() / 1000 / 60;
			    				sb.append("    [" + o.toString() + "] = " + sme.getExpirationTime() + " / expires in " + minutes + " minutes\r\n");
		    				} else {
		    					minutes = new Interval(sme.getExpirationTime(), now).toDurationMillis() / 1000 / 60;
			    				sb.append("    [" + o.toString() + "] = " + sme.getExpirationTime() + " / expired " + minutes + " minutes in the past\r\n");
		    				}
	    				} catch (Exception e) {
		    				sb.append("    [" + o.toString() + "] = " + sme.getExpirationTime() + " / unknown expiration time\r\n");
	    				}
	    			}
	    			sb.append("\r\n");
	    		}
	    	}
	    	
	    	return sb.toString();
    	} catch (Exception e) {
    		log.error("", e);
    		return ("Exception during status check.");
    	}
    }
    
    /** {@inheritDoc} */
    public boolean contains(String partition, Object key) {
        if (partition == null || key == null) {
            return false;
        }

        // check for the special case first
        if (partition.equalsIgnoreCase(sessionPartition) == true) {
        	// look up the data in the session cache
        	return sessionCache.containsKey(key);
        }
        
        if (store.containsKey(partition)) {
            return store.get(partition).containsKey(key);
        }

        return false;
    }

    /**
     * Get an object from the local session store.
     * @param key
     * @return
     */
    private ValueType getLocal(Object key) {
    	log.trace("Searching local storage for key [{}]", key);
    	ValueType vt = sessionCache.get(key);
    	
    	if (vt == null) {
    		log.trace("No object found for key [{}], returning null.", key);
    		return null;
    	}
    	
		// found local copy, return it
		log.trace("Found object for [{}]", key);
		if (vt instanceof SessionID) {
			log.trace("Found SessionID object. Looking for session.");
			SessionID sid = (SessionID) vt;
			// Found our redirection object, get the real thing
			ValueType vt2 = (ValueType) sessionCache.get(sid.sessionID);
			if (vt2 != null) {
				log.trace("Returning Session object from local store for session [{}]", sid.sessionID);
				return vt2;
			}
			// if we do not have a session for the session id, then we need to remove the session id
			log.debug("No session information found for sessionID object [{}]. Removing from cache.", sid);
			this.remove(sessionPartition, (KeyType)key);
			
			return null;
		} else {
			// some other object type
			log.trace("Returning object = [{}]", vt);
			return vt;
		}
    }
    
    /**
     * Get the object from the cluster session store.
     * @param key
     * @return
     */
    @SuppressWarnings("unchecked")
	private ValueType getPersistent(Object key) {
    	log.trace("Searching persistent storage for key [{}]", key);
    	
    	// if we find the object here, this method must put it in the local store
    	// what that means is that the user has moved from one server to another
    	
    	ValueType vt = (ValueType) SessionObjectDAOImpl.getSession(systemName, (String)key);
    	
    	if (vt == null) {
    		log.trace("No object found for key [{}], returning null.", key);
    		return null;
    	}
    	
		// found local copy, return it
		log.trace("Found object for [{}]", key);
		if (vt instanceof SessionID) {
			log.trace("Found SessionID object. Looking for session.");
			SessionID sid = (SessionID) vt;
			// Found our redirection object, get the real thing
	    	ValueType vt2 = (ValueType) SessionObjectDAOImpl.getSession(systemName, sid.sessionID);
			if (vt2 != null) {
				log.trace("Returning Session object from persistent store for session [{}]. Storing object to local storage.", sid.sessionID);
				this.put(sessionPartition, (KeyType)key, vt2);
				return vt2;
			}
			// if we do not have a session for the session id, then we need to remove the session id
			log.debug("No session information found for sessionID object [{}]. Removing from persistent storage.", sid);
			SessionObjectDAOImpl.removeSession(systemName, (String) key);
			
			return null;
		} else {
			// some other object type
			log.trace("Returning object from persistent storage = [{}]. Storing object to local storage.", vt);
			this.put(sessionPartition, (KeyType)key, vt);
			return vt;
		}
    }
    
    /** {@inheritDoc} */
    @SuppressWarnings("unchecked")
	public ValueType get(String partition, Object key) {
        if (partition == null || key == null) {
            return null;
        }

        // check for the special case first
        if (partition.equalsIgnoreCase(sessionPartition) == true) {
        	ValueType vt;
        	// first look for the local copy
        	log.trace("Checking local store for key [{}]", key);
        	vt = getLocal(key);
        	if (vt != null) {
        		log.trace("Returning session object from local store for key [{}]", key);
        		return vt;
        	}
			// did not find object in local store, so look in cluster cache
			// look up the data in the Infinispan cache
			log.trace("Retrieving from persistent storage for key [{}]", key.toString());
			vt = getPersistent(key);
			log.trace("Returning object = [{}]", vt);
			return vt;
        	
        }
        
        if (store.containsKey(partition)) {
        	log.trace("retrieving from [{}] for key [{}]", partition, key.toString());
        	ValueType vt = store.get(partition).get(key);
        	return vt;
        }

        return null;
    }

    /** {@inheritDoc} */
    public Iterator<KeyType> getKeys(String partition) {
        return this.new PartitionEntryIterator(partition);
    }

    /** {@inheritDoc} */
    public Iterator<String> getPartitions() {
    	log.trace("Getting list of partitions.");
    	if (log.isTraceEnabled() == true) {
	    	PartitionIterator pi = this.new PartitionIterator();
	    	while (pi.hasNext() == true) {
	    		log.trace("Partition = [{}]", pi.next());
	    	}
    	}
        return this.new PartitionIterator();
    }

    /** {@inheritDoc} */
    @SuppressWarnings("unchecked")
	public ValueType put(String partition, KeyType key, ValueType value) {
		if (partition == null || key == null) {
			return null;
		}

		// First make sure that the partition exists
		Map<KeyType, ValueType> partitionMap;
		synchronized (store) {
			partitionMap = store.get(partition);
			if (partitionMap == null) {
				partitionMap = new ConcurrentHashMap<KeyType, ValueType>();
				store.put(partition, partitionMap);
			}
		}

		log.debug("Using systemName of [{}] for index [{}]", systemName, key);
		ValueType replacedEntry = null;
		// look for session partition first
		if (partition.equalsIgnoreCase(sessionPartition) == true) {
			log.trace("Session partition requested for storage for key =[{}]", key);
			if (value instanceof SessionManagerEntry) {
				SessionManagerEntry sme = (SessionManagerEntry) value;
				log.trace("Key = [{}] -- Session ID = [{}]", key, sme.getSessionId());
				if (key instanceof String) {
					log.debug("Key is a string = [{}]", key);
					String idx = (String) key;
					if (idx.equals(sme.getSessionId()) == true) {
						// since the index value is the session ID, we need to store the actual object
						log.trace("Storing [{}] in local storage for SessionManagerEntry.", key);
						sessionCache.put(key, value);
						log.trace("Storing [{}] in persistent storage for SessionManagerEntry.", key);
						SessionObjectDAOImpl.putSession(systemName, (String)key, (Object) value);
						
						log.trace("storing to [{}] for key [{}]", partition, key.toString());
						log.trace("value class is [{}]", value.getClass().toString());
					} else {
						// they asked us to store the session under a different index value, but
						// we are going to just store the ID
						SessionID sid = new SessionID();
						sid.sessionID = sme.getSessionId();
						log.trace("Storing [{}] in local storage as SessionID object.", key);
						sessionCache.put(key, (ValueType) sid);
						log.trace("Storing [{}] in persistent storage as SessionID object.", key);
						SessionObjectDAOImpl.putSession(systemName, (String)key, (Object) sid);

						log.trace("storing to [{}] for key [{}]", partition, key.toString());
						log.trace("value is session ID [{}]", sid);
					}
				}
			} else {
				// since this is not an object type we know about, just store it
				log.trace("Storing [{}] in local storage for object [{}].", key, value);
				sessionCache.put(key, value);
				log.trace("Storing [{}] in persistent storage for object [{}].", key, value);
				SessionObjectDAOImpl.putSession(systemName, (String)key, (Object) value);
				
				log.trace("storing to [{}] for key [{}]", partition, key.toString());
				log.trace("value class is [{}]", value.getClass().toString());
			}
		} else {
			// save to local memory only
			log.debug("storing to [{}] for key [{}]", partition, key.toString());
			log.debug("value class is [{}]", value.getClass().toString());
	
			replacedEntry = partitionMap.put(key, value);
		}
		
		appCtx.publishEvent(new AddEntryEvent(this, partition, key, value));

		return replacedEntry;
	}

    /** {@inheritDoc} */
    @SuppressWarnings("unchecked")
	public ValueType remove(String partition, KeyType key) {
        if (partition == null || key == null) {
            return null;
        }

        // look for session partition first
        if (partition.equalsIgnoreCase(sessionPartition) == true) {
        	ValueType removedEntry = (ValueType) sessionCache.remove(key);
	    	log.debug("Removing from local store for key [{}]", key.toString());
	    	// now we must see if we have to remove it from db storage
	    	Object object = SessionObjectDAOImpl.getSession(systemName, (String)key);
	    	if (object != null) {
	    		log.trace("Found object [{}] for key [{}] in persistent storage.", object, key);
	    		if (object instanceof SessionManagerEntry) {
	    			// check and see if this copy has expired before removing it
	    			SessionManagerEntry sme = (SessionManagerEntry) object;
	    			if (sme.isExpired() == true) {
	    				log.trace("Removing key [{}] from persistent storage for SessionManagerEntry.", key);
	    				SessionObjectDAOImpl.removeSession(systemName, (String)key);
	    			} else {
	    				log.trace("Session for key [{}] has not expired in persistent storage. Not removing.", key);
	    			}
	    		} else if (object instanceof SessionID) {
	    			// if we have a SessionID object, we must look for the real session
	    			SessionID sid = (SessionID) object;
	    			// now fetch the session object
	    			Object o2 = SessionObjectDAOImpl.getSession(systemName, sid.sessionID);
	    			if (o2 == null) {
	    				// no session found to match, so remove the SessionID
	    				log.trace("Removing SessionID entry, no matching session object found for key [{}}", key);
	    				SessionObjectDAOImpl.removeSession(systemName, (String) key);
	    			} else {
	    				// we did find the object, does it need to be removed?
		    			// check and see if this copy has expired before removing it
		    			SessionManagerEntry sme = (SessionManagerEntry) o2;
		    			if (sme.isExpired() == true) {
		    				log.trace("Removing key [{}] from persistent storage for SessionID. Base session is expired.", key);
		    				SessionObjectDAOImpl.removeSession(systemName, (String)key);
		    			}
		    			// Note that we leave the actual expired session in place
		    			// so that the storage sweeper service has a consistent view
	    			}
	    		} else {
	    			// not an object we recognize, so just remove it
    				log.trace("Removing key [{}] from persistent storage for object [{}].", key, object);
    				SessionObjectDAOImpl.removeSession(systemName, (String)key);
	    		}
	    	}
	    	
            appCtx.publishEvent(new RemoveEntryEvent(this, partition, key, removedEntry));
            return removedEntry;
        }
        
        if (store.containsKey(partition)) {
            ValueType removedEntry = store.get(partition).remove(key);
	    	log.debug("Removing from  [{}] for key [{}]", partition, key.toString());
            appCtx.publishEvent(new RemoveEntryEvent(this, partition, key, removedEntry));
            return removedEntry;
        }

        return null;
    }

    /** {@inheritDoc} */
    public void setApplicationContext(ApplicationContext ctx) {
        ApplicationContext rootContext = ctx;
        while (rootContext.getParent() != null) {
            rootContext = rootContext.getParent();
        }
        appCtx = rootContext;
    }

    /** An event indicating an item has been added to an storage service. */
    public static class AddEntryEvent<KeyType, ValueType> extends ApplicationEvent {

        /** Serial version UID. */
        private static final long serialVersionUID = -1939512157260059492L;

        /** Storage service to which the item was added. */
        private StorageService<KeyType, ValueType> storageService;

        /** Storage partition to which the item was added. */
        private String partition;

        /** Key to the added item. */
        private KeyType key;

        /** The added item. */
        private ValueType value;

        /**
         * Constructor.
         * 
         * @param storageService storage service to which an item was added
         * @param partition partition to which the entry was added
         * @param key key of the added item
         * @param value added item
         */
        public AddEntryEvent(StorageService<KeyType, ValueType> storageService, String partition, KeyType key,
                ValueType value) {
            super(storageService);
            this.storageService = storageService;
            this.partition = partition;
            this.key = key;
            this.value = value;
        }

        /**
         * Gets the storage service to which an item was added.
         * 
         * @return storage service to which an item was added
         */
        public StorageService<KeyType, ValueType> getStorageService() {
            return storageService;
        }

        /**
         * Gets the partition to which the entry was added.
         * 
         * @return partition to which the entry was added
         */
        public String getPartition() {
            return partition;
        }

        /**
         * Gets the key of the added item.
         * 
         * @return key of the added item
         */
        public KeyType getKey() {
            return key;
        }

        /**
         * Gets the added item.
         * 
         * @return added item
         */
        public ValueType getValue() {
            return value;
        }
    }

    /** An event indicating an item has been removed from an storage service. */
    public static class RemoveEntryEvent<KeyType, ValueType> extends ApplicationEvent {

        /** Serial version UID. */
        private static final long serialVersionUID = 7414605158323325366L;

        /** Storage service to which the item was removed. */
        private StorageService<KeyType, ValueType> storageService;

        /** Storage partition to which the item was removed. */
        private String partition;

        /** Key to the removed item. */
        private KeyType key;

        /** The removed item. */
        private ValueType value;

        /**
         * Constructor.
         * 
         * @param storageService storage service to which an item was removed
         * @param partition partition to which the entry was removed
         * @param key key of the removed item
         * @param value removed item
         */
        public RemoveEntryEvent(StorageService<KeyType, ValueType> storageService, String partition, KeyType key,
                ValueType value) {
            super(storageService);
            this.storageService = storageService;
            this.partition = partition;
            this.key = key;
            this.value = value;
        }

        /**
         * Gets the storage service to which an item was removed.
         * 
         * @return storage service to which an item was removed
         */
        public StorageService<KeyType, ValueType> getStorageService() {
            return storageService;
        }

        /**
         * Gets the partition to which the entry was removed.
         * 
         * @return partition to which the entry was removed
         */
        public String getPartition() {
            return partition;
        }

        /**
         * Gets the key of the removed item.
         * 
         * @return key of the removed item
         */
        public KeyType getKey() {
            return key;
        }

        /**
         * Gets the removed item.
         * 
         * @return removed item
         */
        public ValueType getValue() {
            return value;
        }
    }

    /** An iterator over the partitions of the storage service. */
    public class PartitionIterator implements Iterator<String> {

        /** Iterator over the partitions in the backing store. */
        private Iterator<String> partitionItr;

        /** Current partition. */
        private String currentParition;

        /** Constructor. */
        public PartitionIterator() {
            partitionItr = store.keySet().iterator();
        }

        /** {@inheritDoc} */
        public boolean hasNext() {
            return partitionItr.hasNext();
        }

        /** {@inheritDoc} */
        public String next() {
            currentParition = partitionItr.next();
            return currentParition;
        }

        /** {@inheritDoc} */
        public void remove() {
            Iterator<KeyType> partitionEntries = getKeys(currentParition);
            while (partitionEntries.hasNext()) {
                partitionEntries.next();
                partitionEntries.remove();
            }
            DbStorageService.this.log.debug("Removing partition [{}]", currentParition);
            store.remove(currentParition);
        }
    }

    /** An iterator over the entries of a partition of the storage service. */
    public class PartitionEntryIterator implements Iterator<KeyType> {

        /** Partition on which we are operating. */
        private String partition;

        /** Iterator of keys within the partition. */
        private Iterator<KeyType> keysItr;

        /** Current key within the iteration. */
        private KeyType currentKey;

        /**
         * Constructor.
         * 
         * @param partition partition upon which this iterator operates
         */
        public PartitionEntryIterator(String partition) {
            this.partition = partition;
            // we must be particular about where we get the partition key set
            // if they are asking for the session partition, we must use the infinispan cache
            if (this.partition.equalsIgnoreCase(DbStorageService.sessionPartition) == true) {
            	keysItr = (Iterator<KeyType>) sessionCache.keySet().iterator();
            } else {
            	keysItr = store.get(partition).keySet().iterator();
            }
        }

        /** {@inheritDoc} */
        public boolean hasNext() {
            return keysItr.hasNext();
        }

        /** {@inheritDoc} */
        public KeyType next() {
            currentKey = keysItr.next();
            return currentKey;
        }

        /** {@inheritDoc} */
        public void remove() {
            DbStorageService.this.remove(partition, currentKey);
        }
    }
    
    
}
